import cv2
import numpy as np
import math
from enum import Enum
import time
import onnxruntime

class BODY_TYPE(Enum):
    head = 0
    arm = 1
    body = 2
    leg = 3

BODY_COLOR = {
    BODY_TYPE.head: [0, 255, 0],
    BODY_TYPE.arm: [255, 128, 0],
    BODY_TYPE.body: [255, 51, 255],
    BODY_TYPE.leg: [51, 153, 255]
}
kpt_list = [0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 3, 3, 3, 3, 3, 3]
skeleton = [[16, 14], [14, 12], [17, 15], [15, 13], [12, 13], [6, 12],
            [7, 13], [6, 7], [6, 8], [7, 9], [8, 10], [9, 11], [2, 3],
            [1, 2], [1, 3], [2, 4], [3, 5]]
skeleton_list = [3, 3, 3, 3, 2, 2, 2, 2, 1, 1, 1, 1, 0, 0, 0, 0, 0]

class Recognition:
	def __init__(self, model_path, gpu_enabled=True, color_all=True):
		self.model_path = model_path
		self.color_all = color_all
		self.running = True
		self.point = (0, 0)
		self.origin = None			# 原图像
		self.deal_pic = None		# 处理线上的图像
		self.recog_result = None	# 处理线上的结果
		self.result = None			# 识别分类结果
		if gpu_enabled:
			self.session = onnxruntime.InferenceSession(self.model_path, providers=['CUDAExecutionProvider'])
		else:
			self.session = onnxruntime.InferenceSession(self.model_path)
		self.input_name = self.session.get_inputs()[0].name
	
	def GetResult(self, img):
		self.origin = img
		self.result = None
		self.deal_pic = None
		self.recog_result = None
		self.PreProcess()
		self.ModelInference()
		self.Nms()
		if self.recog_result.size == 0:
			return self.result, img
		else:
			self.PostProcess()
			result = cv2.resize(self.deal_pic, img.shape[1::-1])
			return self.result, result
		
	def PreProcess(self):
		self.origin = cv2.resize(self.origin, (640, 640), interpolation=cv2.INTER_LINEAR)
		self.deal_pic = (self.origin - 127.5) / 127.5
		self.deal_pic = np.asarray(self.deal_pic, dtype=np.float32)
		self.deal_pic = np.expand_dims(self.deal_pic, axis=0)
		self.deal_pic = np.transpose(self.deal_pic, (0, 3, 1, 2))
	
	def ModelInference(self, conf = 0.5):
		self.recog_result = self.session.run([], {self.input_name: self.deal_pic})[0]
		self.recog_result = np.transpose(self.recog_result[0], (1, 0))
		self.recog_result = self.recog_result[self.recog_result[:, 4] > conf]
		# (中心点xy, 宽高wh) -> (左上xy, 右下xy)
		temp = np.copy(self.recog_result)
		self.recog_result[:, 0] = temp[:, 0] - temp[:, 2] / 2
		self.recog_result[:, 1] = temp[:, 1] - temp[:, 3] / 2
		self.recog_result[:, 2] = temp[:, 0] + temp[:, 2] / 2
		self.recog_result[:, 3] = temp[:, 1] + temp[:, 3] / 2

	def Nms(self, iou_thresh=0.5):
		x1 = self.recog_result[:, 0]
		y1 = self.recog_result[:, 1]
		x2 = self.recog_result[:, 2]
		y2 = self.recog_result[:, 3]
		scores = self.recog_result[:, 4]
		areas = (x2 - x1 + 1) * (y2 - y1 + 1)
		order = scores.argsort()[::-1]
		keep = []
		while order.size > 0:
			i = order[0]
			keep.append(i)
			xx1 = np.maximum(x1[i], x1[order[1:]])
			yy1 = np.maximum(y1[i], y1[order[1:]])
			xx2 = np.minimum(x2[i], x2[order[1:]])
			yy2 = np.minimum(y2[i], y2[order[1:]])
			w = np.maximum(0.0, xx2 - xx1 + 1)
			h = np.maximum(0.0, yy2 - yy1 + 1)
			inter = w * h
			ovr = inter / (areas[i] + areas[order[1:]] - inter)
			inds = np.where(ovr <= iou_thresh)[0]
			order = order[inds + 1]
		output = []
		for i in keep:
			output.append(self.recog_result[i].tolist())
		self.recog_result = np.array(output)

	def PostProcess(self):
		bboxes, scores, kpts = self.recog_result[:, 0:4], self.recog_result[:, 4], self.recog_result[:, 5:]
		self.deal_pic = self.origin.copy()
		distance, id = -1, -1
		for idx in range(len(bboxes)):
			bbox, score, kpt = bboxes[idx], scores[idx], kpts[idx]
			# 人体框选染色
			if self.color_all:
				color = (0, 0, 255)
				cv2.rectangle(self.deal_pic, (int(bbox[0]), int(bbox[1])), (int(bbox[2]), int(bbox[3])), color, 2)
				cv2.putText(self.deal_pic, 'score: {:.2f}'.format(score), (int(bbox[0] + 5), int(bbox[1]) + 15), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
			_distance, _id = self.ProcessSkeletonKpts(kpt)
			if distance < 0 or (_id >= 0 and _distance < distance):
				distance, id = _distance, _id
		if id >= 0:
			self.result = BODY_TYPE(skeleton_list[id])

	def ProcessSkeletonKpts(self, kpts, threshold = 0.3):
		steps = 3
		min_val, max_val = -1, 150
		f_pos1, f_pos2 = (-1, -1), (-1, -1)
		f_sk_id = -1
		# 染色目标节点和骨骼节点
		if self.color_all:
			cv2.circle(self.deal_pic, self.point, 3, (255, 255, 255), -1)
			for kid in range(len(kpts) // steps):
				r, g, b = BODY_COLOR[BODY_TYPE(kpt_list[kid])]
				x_coord, y_coord = kpts[steps * kid], kpts[steps * kid + 1]
				conf = kpts[steps * kid + 2]
				if conf > threshold:
					cv2.circle(self.deal_pic, (int(x_coord), int(y_coord)), 3, (int(r), int(g), int(b)), -1)
		for sk_id, sk in enumerate(skeleton):
			pos1 = (int(kpts[(sk[0] - 1) * steps]), int(kpts[(sk[0] - 1) * steps + 1]))
			pos2 = (int(kpts[(sk[1] - 1) * steps]), int(kpts[(sk[1] - 1) * steps + 1]))
			conf1 = kpts[(sk[0] - 1) * steps + 2]
			conf2 = kpts[(sk[1] - 1) * steps + 2]
			if conf1 <= threshold or conf2 <= threshold:
				continue
			distance = 0
			v1 = np.array([self.point[0] - pos1[0], self.point[1] - pos1[1]])
			v2 = np.array([self.point[0] - pos2[0], self.point[1] - pos2[1]])
			v3 = np.array([pos1[0] - pos2[0], pos1[1] - pos2[1]])
			if v1.dot(v3) >= 0:
				distance = math.sqrt(v1.dot(v1))
			elif v2.dot(v3) <= 0:
				distance = math.sqrt(v2.dot(v2))
			else:
				rotate90 = np.array([[0, 1], [-1, 0]])
				v3 = v3.dot(rotate90)
				distance = abs(v1.dot(v3)) / math.sqrt(v3.dot(v3))
			if distance <= max_val and (distance < min_val or min_val < 0):
				min_val = distance
				f_pos1, f_pos2 = pos1, pos2
				f_sk_id = sk_id
			if self.color_all:
				r, g, b = BODY_COLOR[BODY_TYPE(skeleton_list[sk_id])]
				# cv2.putText(self.deal_pic, str(sk_id), ((pos1[0] + pos2[0]) // 2, (pos1[1] + pos2[1]) // 2), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (int(r), int(g), int(b)), 1)
				cv2.line(self.deal_pic, pos1, pos2, (int(r), int(g), int(b)), thickness=1)
		# 处理距离最近的骨骼
		if min_val >= 0 and self.color_all:
			cv2.line(self.deal_pic, f_pos1, f_pos2, (255, 255, 255), thickness=2)
		return min_val, f_sk_id